from vnpy.trader.optimize import OptimizationSetting
from vnpy_ctastrategy.backtesting import BacktestingEngine
from vnpy.trader.constant import Interval
from vnpy.trader.object import TickData, BarData, TradeData, OrderData
from vnpy.trader.utility import BarGenerator, ArrayManager
import pandas as pd
from datetime import datetime
from vnpy.trader.setting import SETTINGS

from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,
)

class MomentumReversalStrategy(CtaTemplate):
    """
    量价共振短线策略
    """
    author = "yzj"
    # 配置数据库
    SETTINGS["database.name"] = "mysql"              # 可以根据自己的需求选择数据库，这里使用的是TDengine
    SETTINGS["database.database"] = "vnpy"
    SETTINGS["database.host"] = "127.0.0.1"
    SETTINGS["database.port"] = 3306
    SETTINGS["database.user"] = "root"
    SETTINGS["database.password"] = "12345678"
    # 策略参数
    holding_days = 5  # 持有天数
    stop_loss = 0.03  # 止损比例
    take_profit = 0.05  # 止盈比例
    position_ratio = 0.2  # 单票仓位比例
    
    # 策略变量
    entry_price = 0
    holding_period = 0
    current_trade_date = None
    
    parameters = ["holding_days", "stop_loss", "take_profit", "position_ratio"]
    variables = ["entry_price", "holding_period"]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        
        self.bg = BarGenerator(self.on_bar)
        self.am = ArrayManager(size=100)

    def on_init(self):
        """
        初始化策略
        """
        self.write_log("策略初始化")
        self.load_bar(10)  # 加载10天数据

    def on_start(self):
        """
        策略启动
        """
        self.write_log("策略启动")

    def on_stop(self):
        """
        策略停止
        """
        self.write_log("策略停止")

    def on_tick(self, tick: TickData):
        """
        Tick数据更新
        """
        self.bg.update_tick(tick)

    def on_bar(self, bar: BarData):
        """
        K线数据更新
        """
        self.am.update_bar(bar)
        if not self.am.inited:
            return
        
        self.current_trade_date = bar.datetime.date()
        
        # 如果有持仓，检查止损止盈和平仓条件
        if self.pos != 0:
            self.holding_period += 1
            
            # 计算当前盈亏
            if self.pos > 0:
                profit_ratio = (bar.close_price - self.entry_price) / self.entry_price
            else:
                profit_ratio = (self.entry_price - bar.close_price) / self.entry_price
            
            # 止损检查
            if profit_ratio <= -self.stop_loss:
                self.close_all()
                return
            
            # 止盈检查
            if profit_ratio >= self.take_profit:
                self.close_all()
                return
            
            # 持有期结束平仓
            if self.holding_period >= self.holding_days:
                self.close_all()
                return
        
        # 如果没有持仓，检查买入信号
        else:
            if self.check_buy_signal(bar):
                # 计算买入数量
                size = int(self.cta_engine.capital * self.position_ratio / bar.close_price / 100) * 100
                if size > 0:
                    self.buy(bar.close_price, size)
                    self.entry_price = bar.close_price
                    self.holding_period = 0

    def check_buy_signal(self, bar: BarData) -> bool:
        """
        检查买入信号
        这里需要接入之前写的选股逻辑
        """
        # 在实际使用中，这里应该调用选股函数
        # 为了演示，我们使用简化版的信号判断
        
        if len(self.am.close) < 30:
            return False
        
        # 简化版信号条件（实际应该使用完整的选股逻辑）
        ret_5 = (bar.close_price / self.am.close[-5] - 1)
        vol_ma5 = self.am.volume[-5:].mean()
        ma20 = self.am.close[-20:].mean()
        
        condition1 = ret_5 < -0.08
        condition2 = bar.volume > vol_ma5 * 0.8
        condition3 = bar.volume > self.am.volume[-2] * 1.3
        condition4 = abs(bar.close_price - ma20) / ma20 < 0.02
        
        return condition1 and condition2 and condition3 and condition4

    def close_all(self):
        """
        平仓所有头寸
        """
        if self.pos > 0:
            self.sell(self.last_bar.close_price, abs(self.pos))
        elif self.pos < 0:
            self.cover(self.last_bar.close_price, abs(self.pos))
    
    @classmethod
    def generate_settings(cls, optimization_target="sharpe_ratio"):
        """
        生成参数优化设置
        
        Args:
            optimization_target: 优化目标，可选值:
                - "sharpe_ratio": 夏普比率
                - "total_return": 总收益率
                - "max_drawdown": 最大回撤（负优化）
                - "profit_factor": 盈亏比
        """
        setting = OptimizationSetting()
        
        # 设置优化目标
        if optimization_target == "max_drawdown":
            # 最大回撤需要负优化（越小越好）
            setting.set_target("max_drawdown")
            setting.set_direction("minimize")
        else:
            setting.set_target(optimization_target)
        
        # 添加优化参数范围
        # 1. 持仓参数
        setting.add_parameter("holding_days", 3, 10, 1)          # 持有天数：3-10天
        setting.add_parameter("position_ratio", 0.1, 0.5, 0.05)  # 仓位比例：10%-50%
        
        # 2. 风险控制参数
        setting.add_parameter("stop_loss", 0.01, 0.05, 0.005)    # 止损比例：1%-5%
        setting.add_parameter("take_profit", 0.03, 0.08, 0.01)   # 止盈比例：3%-8%
        
        # 3. 信号参数（新增）
        setting.add_parameter("volume_multiplier", 1.2, 2.0, 0.1)        # 成交量倍数：1.2-2.0
        setting.add_parameter("retracement_threshold", 0.05, 0.12, 0.01) # 回撤阈值：5%-12%
        setting.add_parameter("ma_period", 10, 30, 5)                    # 均线周期：10-30
        setting.add_parameter("ma_deviation", 0.01, 0.04, 0.005)         # 均线偏离：1%-4%
        
        return setting


def run_backtest():
    """
    运行回测
    """
    # 创建回测引擎
    engine = BacktestingEngine()
    
    # 设置回测参数
    engine.set_parameters(
        vt_symbol="832522.BSE",  # 股票代码
        interval=Interval.DAILY,
        start=datetime(2025, 1, 1),
        end=datetime(2025, 8, 15),
        rate=0.0003,  # 手续费
        slippage=0.001,  # 滑点
        size=1,  # 合约乘数
        pricetick=0.01,  # 价格跳动
        capital=1000000,  # 初始资金
    )
    
    # 添加策略
    engine.add_strategy(MomentumReversalStrategy, {})
    
    # 加载数据
    engine.load_data()
    
    # 运行回测
    engine.run_backtesting()
    
    # 计算回测结果
    df = engine.calculate_result()
    statistics = engine.calculate_statistics()
    
    # 输出结果
    print("回测结果:")
    print(df.tail())
    print("\n统计指标:")
    for key, value in statistics.items():
        print(f"{key}: {value}")
    
    # 绘制图表
    engine.show_chart()

if __name__ == "__main__":
    run_backtest()
